import os
import pysftp
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Union

# 定义匹配规则
MATCH_SUFFIX = '_rgba'

MAX_RETRIES = 3  # 最大重试次数

def sftp_download_file(cnopts: pysftp.CnOpts, remote_file: str, local_file: str, retries: int = 0) -> bool:
    try:
        with pysftp.Connection(
            host='172.18.0.83',
            username='root',
            password='rabbit1023../',
            cnopts=cnopts
        ) as sftp:
            sftp.get(remote_file, local_file)
            print(f'下载完成: {remote_file}')
            return True
    except Exception as e:
        if retries < MAX_RETRIES:
            print(f'download exception: {e}, 重试 ({retries + 1}/{MAX_RETRIES})...')
            return sftp_download_file(cnopts, remote_file, local_file, retries + 1)
        else:
            print(f'download exception: {e}')
            return False

def is_dir(sftp: pysftp.Connection, path: str) -> bool:
    try:
        return sftp.isdir(path)
    except Exception as e:
        print(f'is_dir exception: {e}')
        return False

def is_file(sftp: pysftp.Connection, path: str) -> bool:
    try:
        return sftp.isfile(path)
    except Exception as e:
        print(f'is_file exception: {e}')
        return False

def sftp_upload_file(cnopts: pysftp.CnOpts, local_file: str, remote_file: str, retries: int = 0) -> bool:
    try:
        with pysftp.Connection(
            host='172.18.0.83',
            username='root',
            password='rabbit1023../',
            cnopts=cnopts
        ) as sftp:
            sftp.put(local_file, remote_file)
            print(f'上传完成: {local_file}')
            return True
    except Exception as e:
        if retries < MAX_RETRIES:
            print(f'upload exception: {e}, 重试 ({retries + 1}/{MAX_RETRIES})...')
            return sftp_upload_file(cnopts, local_file, remote_file, retries + 1)
        else:
            print(f'upload exception: {e}')
            return False

def main():
    host = '172.18.0.83'
    port = 22
    username = 'root'
    password = 'rabbit1023../'
    local = r'C:\Users\admin\Desktop\1'  # 注意末尾的斜杠
    remote = "/nfs/data/data/2-6-1/"

    cnopts = pysftp.CnOpts()
    cnopts.hostkeys = None

    with pysftp.Connection(
        host=host,
        username=username,
        password=password,
        cnopts=cnopts
    ) as sftp:
        print('链接成功')

        try:
            if os.path.isdir(local):  # 判断本地参数是目录还是文件
                if is_dir(sftp, remote):  # 检查远程路径是否为目录
                    tasks = []
                    with ThreadPoolExecutor(max_workers=5) as executor:
                        for f in os.listdir(local):  # 遍历本地目录
                            local_path = os.path.join(local, f)
                            remote_path = os.path.join(remote, f)

                            if os.path.isfile(local_path) and MATCH_SUFFIX in f:
                                future = executor.submit(sftp_upload_file, cnopts, local_path, remote_path)
                                tasks.append(future)
                        for task in tasks:
                            result = task.result()  # 等待所有任务完成
                            if not result:
                                print(f'上传失败: {task}')
                else:
                    print(f'远程路径 {remote} 不是目录')
            else:
                sftp.put(local, remote)
                print('上传完成')
        except Exception as e:
            print(f'upload exception: {e}')

if __name__ == '__main__':
    main()
